# websocket/manager.py
from typing import List, Dict, Optional
from fastapi import WebSocket


class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}  # user_id -> WebSocket
        self.room_members: Dict[str, Dict[str, WebSocket]] = {}  # room_name -> {user_id: WebSocket}

    async def connect(self, websocket: WebSocket, user_id: Optional[str] = None):
        await websocket.accept()
        if user_id:
            self.active_connections[user_id] = websocket
        return websocket

    def disconnect(self, user_id: str):
        if user_id in self.active_connections:
            del self.active_connections[user_id]

    def join_room(self, room_name: str, user_id: str, websocket: WebSocket):
        if room_name not in self.room_members:
            self.room_members[room_name] = {}
        self.room_members[room_name][user_id] = websocket

    def leave_room(self, room_name: str, user_id: str):
        if room_name in self.room_members and user_id in self.room_members[room_name]:
            del self.room_members[room_name][user_id]
            if not self.room_members[room_name]:
                del self.room_members[room_name]

    async def send_personal_message(self, message: dict, websocket: WebSocket):
        await websocket.send_json(message)

    async def broadcast_to_room(self, room_name: str, message: dict):
        if room_name in self.room_members:
            for connection in self.room_members[room_name].values():
                await connection.send_json(message)

    def get_clients_in_room(self, room_name: str) -> list:
        if room_name in self.room_members:
            return list(self.room_members[room_name].keys())
        return []